package assignment01;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.Set;
import java.util.TreeSet;


import nlp.langmodel.LanguageModel;
import nlp.util.Counter;
import nlp.util.CounterMap;
import nlp.util.Pair;

public class TrigramModel implements LanguageModel {
	private Collection<List<String>> trainingCollection;
	private Collection<List<String>> validationCollection;
//	private CounterMap<String, Counter<String>> trigramCounterMap = new CounterMap<String, Counter<String>>();
	private CounterMap<Bigram<String, String>, String> trigramCounterMap = new CounterMap<Bigram<String,String>, String>();
	private Counter<Integer> nCounter = new Counter<Integer>();
	private Counter<Bigram<String, String>> historyCounter = new Counter<Bigram<String, String>>();
	private Counter<String> wordCounter = new Counter<String>();
	private static final double INCREMENT_FACTOR = 1.0;
	private final static Integer THRESHOLD = 5;
	private Double[] N = new Double[THRESHOLD + 2];
	private Double[] C = new Double[THRESHOLD + 1];
	private static final String START = "<S>";
	private static final String STOP = "</S>";
	private double vocabulary = 0.0;
	private double unseenCount = 0.0;
	private double seenCount = 0.0;
	
	public TrigramModel(Collection<List<String>> trainingCollection, Collection<List<String>> validationCollection) {
		this.trainingCollection = trainingCollection;
		this.validationCollection = validationCollection;
		train();
		this.vocabulary = wordCounter.keySet().size();
		this.seenCount = trigramCounterMap.totalSize();
		this.unseenCount = Math.pow(vocabulary, 3.0) - this.seenCount;
			
		System.out.println("vacabulary size: " + vocabulary);
		System.out.println("before smoothing, total seen type number is " + seenCount);
		System.out.println("before smoothing, total unseen type number is " + unseenCount);
		System.out.println("before smoothing, total seen token number is " + trigramCounterMap.totalCount());
		smooth();
		System.out.println("C[0] is " + C[0]);
		System.out.println("after smoothing, total token is " + validate() + ", and N[1] is " + N[1]);
		solveHistory();
	}
	
	void train() {
		for (List<String> sentence : trainingCollection) {
			List<String> stopSentence = new ArrayList<String>(sentence);
			stopSentence.add(STOP);

			for (String word : stopSentence) {
				wordCounter.incrementCount(word, INCREMENT_FACTOR);
			}
		}
		
		for (List<String> sentence : trainingCollection) {
			List<String> stopSentence = new ArrayList<String>(sentence);
			stopSentence.add(0, START);
			stopSentence.add(0, START);
			stopSentence.add(STOP);
			
			Iterator<String> iter = stopSentence.iterator();
			String first = iter.next();
			String second = iter.next();
			String third = iter.next();

			while (iter.hasNext() || !third.equals(STOP)) {
				Bigram<String, String> bigram = new Bigram<String, String>(first, second);
				Counter<String> counter = trigramCounterMap.getCounter(bigram);
				counter.incrementCount(third, INCREMENT_FACTOR);
				
				first = second;
				second = third;
				if (iter.hasNext()) {
					third = iter.next();
				}
			}
		}
		
		// find count of counts
		for (Bigram<String, String> history : trigramCounterMap.keySet()) {
			Counter<String> counter = trigramCounterMap.getCounter(history);
			for (String third : counter.keySet()) {
				double count = counter.getCount(third);
				nCounter.incrementCount((int)count, INCREMENT_FACTOR);
			}
		}
	}
	
	void smooth() {
		Set<Integer> treeSet = new TreeSet<Integer>();
		treeSet.addAll(nCounter.keySet());
		Iterator<Integer> iter = treeSet.iterator();
		
		int next = 0;
		while (iter.hasNext() && next <= THRESHOLD) {
			next = iter.next();
			N[next] = nCounter.getCount(next);
		}
		N[0] = unseenCount;
		
		for(int c = 1; c <= THRESHOLD; c++) {
			double cStar = Math.log((c + 1) * N[c + 1] / N[c]
					- (c * (THRESHOLD + 1) * N[THRESHOLD + 1]) / N[1])
					- Math.log(1.0 - (THRESHOLD + 1) * N[THRESHOLD + 1] / N[1]);
			cStar = Math.pow(Math.E, cStar);
			C[c] = cStar;
		}
		C[0] = N[1] / unseenCount;
	}
	

	double validate() {
		double sum = 0.0;
		for (Bigram<String, String> history : trigramCounterMap.keySet()) {
			Counter<String> counter = trigramCounterMap.getCounter(history);
			for (String third : counter.keySet()) {
				double count = counter.getCount(third);
				if ((int)count <= THRESHOLD) {
					sum += C[(int)count];
				} else {
					sum += count;
				}
			}
		}
		return sum;
	}
	
	void solveHistory() {
		for (Bigram<String, String> history : trigramCounterMap.keySet()) {
			Counter<String> counter = trigramCounterMap.getCounter(history);
			double t = vocabulary + 1;
			double size = counter.size();
			double sum = 0.0;
			for (String second : counter.keySet()) {
				double count = counter.getCount(second);
				if ((int)count <= THRESHOLD) {
					sum += C[(int)count];
				} else {
					sum += count;
				}
			}
//			double marginal = (t - size) * C[0] + sum;
			double marginal = sum - 1.0 / (t - 1.0);
			historyCounter.incrementCount(history, marginal);
		}
	}
	
	double getWordLogProbability(List<String> sentence, int index) {
		if (index <= 1) {
			try {
				throw new Exception();
			} catch(Exception e) {
				System.err.println("Index out of range!");
			}
		}
		String first = sentence.get(index - 2);
		String second = sentence.get(index - 1);
		Bigram<String, String> history = new Bigram<String, String>(first, second);
		String word = sentence.get(index);
		
		// if history unseen
		return getPWord(history, word);
	}

	double getPWord(Bigram<String, String> history, String word) {
		if (!trigramCounterMap.containsKey(history)) {
			return Math.log(1.0 / (vocabulary + 1.0)) / Math.log(2.0);
		}
		
		double count = 0.0;
		Counter<String> counter = trigramCounterMap.getCounter(history);
		int intCount = (int) counter.getCount(word);
		if (intCount <= THRESHOLD) {
//			if (intCount == 0) {
//				System.out.println("unknown word: " + word);
//			}
			count = C[intCount];
		} else {
			count = intCount;
		}
		double marginal = historyCounter.getCount(history);
//		if (Double.isInfinite(Math.log(count / marginal) / Math.log(2.0))) {
//			System.out.println("shit!! count is " + count + ", and marginal is " + marginal);
//		}
//		System.out.println("word probability is " + Math.log(count / marginal) / Math.log(2.0));
		return Math.log(count / marginal) / Math.log(2.0);
	}
	
	
	@Override
	public double getSentenceProbability(List<String> sentence) {
		List<String> stoppedSentence = new ArrayList<String>(sentence);
		stoppedSentence.add(0, START);
		stoppedSentence.add(0, START);
		stoppedSentence.add(STOP);
		double probability = 0.0;
		for (int i = 2; i < stoppedSentence.size(); i++) {
			probability += getWordLogProbability(stoppedSentence, i);
		}
//		System.out.println("sentence p is " + probability);
//		if (Double.isInfinite(probability) || Double.isNaN(probability)) {
//			System.out.println("shit!!");
//		}
		return Math.pow(2.0, probability);
	}
	
	String generateWord(Bigram<String, String> history) {
		double sample = Math.random();
		double sum = 0.0;
		Counter<String> counter = trigramCounterMap.getCounter(history);
		double marginal = historyCounter.getCount(history);
		for (String word : counter.keySet()) {
			int intCount = (int) counter.getCount(word);
			if (intCount <= THRESHOLD) {
				sum += C[intCount];
			} else {
				sum += intCount;
			}
			if (sum / marginal >= sample) {
//				System.out.println(word);
				return word;
			}
//			System.out.println("sum / marginal is " + sum / marginal);
		}
//		return "*UNKNOWN*";
		return STOP;
	}
	
	@Override
	public List<String> generateSentence() {
		List<String> sentence = new ArrayList<String>();
		String first = START;
		Bigram<String, String> history = new Bigram<String, String>(START, START);
		String word = generateWord(history);
		int i = 0;
		while (!word.equals(STOP)) {
			sentence.add(word);
			Bigram<String, String> hist = new Bigram<String, String>(first, word);
			first = word;
//			System.out.println("hist is " + hist);
			word = generateWord(hist);
			i++;
		}
		return sentence;
	}
	
	
	/**
	 // test for probability sum add up to 1
	String generateWord(Bigram<String, String> history) {
		Counter<String> counter = trigramCounterMap.getCounter(history);
		double sum = 0.0;
		double marginal = historyCounter.getCount(history);
		for (String word : counter.keySet()) {
			int intCount = (int) counter.getCount(word);
			if (intCount <= THRESHOLD) {
				sum += C[intCount];
			} else {
				sum += intCount;
			}
		}
		System.out.println("sum/marginal + 1/t is " + (sum / marginal + 1.0 / (this.vocabulary + 1.0)));
		return null;
	}
	
	@Override
	public List<String> generateSentence() {
		// TODO Auto-generated method stub
		int i = 0;
		for (Bigram<String, String> history : trigramCounterMap.keySet()) {
			generateWord(history);
			i++;
			if (i > 10) {
				break;
			}
		}
		return null;
	}
	*/
}
